#include "Graph.h"
#include "solver.h"
#include <map>
#include <vector>

using namespace std;


void partial_reverse_cuthill_mckee(const int n, const int *row, const int *col, int *perm, int threshold);

void get_csc_degrees(int n, int *row, int *col, int *degrees) {
    for (int j = 0; j < n; j++) {
      for (int i = col[j]; i < col[j + 1]; i++) {
        degrees[j]++;
        degrees[row[i]]++;
      }
    }
  }


int get_thresholds(int n, int *degrees, int len, int *thresholds, int *a) {
    std::map<int, int> thres;
    std::map<int, int> cnt;
    for (int i = 0; i < n; i++) {
        cnt[ degrees[i] ]++;
    }


    // 等效带宽公式： B = alpha * t + S
    // 选择 len 个alpha，计算对应的B，选择最小B对应的t，t即阈值
    int threshold;
    for (int i = 0; i < len; i++) {
        int alpha = 1 << i; // 参数alpha
        int S = n;
        int min_cost = INT_MAX;
        threshold = 0;
        // 初始所有节点均处于S， 阈值t从左向右滑动，
        // 和threshold2完全等价
        for (const auto &kvp : cnt) { 
            S -= kvp.second;          
            unsigned long long cost = alpha * kvp.first + S;
            if (cost < min_cost) {
                min_cost = cost;
                threshold = kvp.first;
            }
        }
        // thres.insert(threshold);
        thres[threshold] = alpha;  // 记录当前alpha对应的t，可能会有重复的threshold
    }
    int i = 0;
    for (const auto &k : thres) {
        thresholds[i] = k.first;  // t
        a[i] = k.second;  // alpha
        i++;
    }
    return i;
}

void get_beta(int n, int *row, int *col, int *perm, int *degrees, int threshold, int *beta, int *s) {
    *s = 0;
    for (int i = 0; i < n; i++) {
        *s += degrees[i] > threshold;
    }

	int *perm_inv = (int *)malloc(n * sizeof(int));
	for (int i = 0; i < n; ++i) perm_inv[ perm[i] ] = i; // 反函数

    *beta = 0;
    for (int j = 0; j < n; j++) {
        if (degrees[j] > threshold) continue;
        for (int i = col[j]; i < col[j + 1]; i++) {
            if (degrees[ row[i] ] > threshold) continue;
            int b = perm_inv[ row[i] ] - perm_inv[j];
            b = b < 0 ? 0 - b : b;
            *beta = MAX(b, *beta);
        }
    }
    free(perm_inv);
}

extern "C" {
    #ifndef HSL_MC64D_H
    #define HSL_MC64D_H
    
    
    #ifndef mc64_default_control
    #define mc64_control mc64_control_d
    #define mc64_info mc64_info_d
    #define mc64_default_control mc64_default_control_d
    #define mc64_matching mc64_matching_d
    #endif
    
    typedef double mc64pkgtype_d_;
    
    struct mc64_control_d {
       int f_arrays;
       int lp;
       int wp;
       int sp;
       int ldiag;
       int checking;
    };
    
    struct mc64_info_d {
       int flag;
       int more;
       int strucrank;
       int stat;
    };
    
    /* Set default values of control */
    void mc64_default_control_d(struct mc64_control *control);
    /* Find a matching, and (optionally) scaling */
    void mc64_matching_d(int job, int matrix_type, int m, int n, const int *ptr,
       const int *row, const mc64pkgtype_d_ *cval,
       const struct mc64_control *control,
       struct mc64_info *info, int *perm, mc64pkgtype_d_ *scale);
    #endif
}


void cperm(int* perm, MATRIX* A){
    int nonz = A->real_nonz;
    int n = A->real_n;

    int *row2 = (int *) malloc(nonz * sizeof(int));
    int *col2 = (int *) malloc((n + 1) * sizeof(int));
    double *a2   = (double *) malloc(nonz * sizeof(double));


    col2[0] = 0;
    for (int j = 0; j < n; j++) {
        col2[j + 1] = A->col[ perm[j] + 1 ] - A->col[ perm[j] ];
    }
    for (int j = 0; j < n; j++) {
        col2[j + 1] += col2[j];
    }


    for (int j = 0; j < n; j++) {
        int i1 = A->col [ perm[j] ];
        for (int i2 = col2[j]; i2 < col2[j + 1]; i2++, i1++) {
            row2[i2] = A->row[i1];
            a2  [i2] = A->a[i1];
        }
    }

    free(A->row);
    free(A->col);
    free(A->a);
    A->row = row2;
    A->col = col2;
    A->a = a2;

}


void rperm(int* perm, MATRIX* A){
    int nonz = A->real_nonz;
    int n = A->real_n;

    int *row2 = (int *) malloc(nonz * sizeof(int));
    for (int j = 0; j < n; j++) {
        for (int i = A->col[j]; i < A->col[j + 1]; ++i) {
            row2[i] = perm[ A->row[i] ];
        }
    }

    free(A->row);
    A->row = row2;
}


int get_bw(MATRIX *A, int n)
{
    int bw = 0;
    int *row = A->row;
    int *col = A->col;
    for (int i = 0; i < n; ++i)
        for (int j = row[i]; j < row[i + 1] && col[j] < n; ++j)
            bw = MAX(bw, ABS(i - col[j]));
    return bw;
}


void get_mc64(MATRIX* A){
    if(A->type != CSC){
        printf("MC64 reorder must be CSC format.\n");
        return;
    }

    auto row = A->row;
    auto col = A->col;
    auto a = A->a;
    int m = A->real_m;
    int n = A->real_n;

    int *perm_mc64 = (int *)malloc(2 * n * sizeof(int));
    int *perm_mc64_inv = perm_mc64 + n;

    mc64_control control;
    mc64_info info;

    mc64_default_control(&control);

    int job = 5;
    int matrix_type = 2;


    // col: ptr ; length=n+1
    // row: index ; length=ptr[n]
    mc64_matching(job, matrix_type, m, n, col, row, a, &control, &info, perm_mc64, NULL);
    if (info.flag < 0) {
        printf("Failure of mc64_matching with info.flag=%d\n", info.flag);
        ABORT("STOP");
    }

    for (int i = 0; i < n; i++) perm_mc64[perm_mc64[i + n]] = i;

    cperm(perm_mc64, A);
}


void get_rcm(MATRIX* A){
    if(A->type != CSC){
        printf("RCM reorder must be CSC format.\n");
        return;
    }

    int* perm = (int *)malloc(A->real_n * sizeof(int));
    int* perm_inv = (int *)malloc(A->real_n * sizeof(int));
    
    int threshold = INT_MAX;
    partial_reverse_cuthill_mckee(A->real_n, A->row, A->col, perm, threshold);

    for(int i = 0; i < A->real_n; ++i){
        perm_inv[perm[i]] = i;
    }

    cperm(perm, A);
    rperm(perm_inv, A);

    free(perm);
    free(perm_inv);
}


void get_prcm(MATRIX* A){
    if(A->type != CSC){
        printf("PRCM reorder must be CSC format.\n");
        return;
    }
    int* degrees = (int *)malloc(A->real_n * sizeof(int));
    get_csc_degrees(A->real_n, A->row, A->col, degrees);
    int len = 15;
    int threshold[len];
    int a[len];
    len = get_thresholds(A->real_n, degrees, len-1, threshold, a);
    threshold[len++] = INT_MAX;  // 最后一个阈值为无穷大

    int betas[len], ss[len], *perms[len];

    #pragma omp parallel for num_threads(std::min(len, omp_get_max_threads()))
    for (int i = 0; i < len; i++) {
        perms[i] = (int *)malloc(A->real_n * sizeof(int));
        partial_reverse_cuthill_mckee(A->real_n, A->row, A->col, perms[i], threshold[i]);
        get_beta(A->real_n, A->row, A->col, perms[i], degrees, threshold[i], &betas[i], &ss[i]);
    }

    free(degrees);

    int index = 0;
    int min = betas[0] + ss[0]; 
    for (int i = 1; i < len; i++) {
        if (betas[i] + ss[i] < min) {
            min = betas[i] + ss[i];
            index = i;
        }
    }

    int *perm = (int *)malloc(A->real_n * sizeof(int));
    for (int i = 0; i < A->real_n; i++) {
        perm[i] = perms[index][i];
    }
    int *perm_inv = (int *)malloc(A->real_n * sizeof(int));
    for (int i = 0; i < A->real_n; i++) {
        perm_inv[perm[i]] = i;
    }
    

    cperm(perm, A);
    rperm(perm_inv, A);
    for (int i = 0; i < len; i++) {
        free(perms[i]);
    }
    
}
